package 链表题目;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

//查找和最小的K对数字
//https://leetcode.cn/problems/find-k-pairs-with-smallest-sums/description/

//给定两个以 非递减顺序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k 。
//定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。
//请找到和最小的 k 个数对 (u1,v1),  (u2,v2)  ...  (uk,vk) 。
class 查找和最小的K对数字 {
    // 我写的
    public List<List<Integer>> kSmallestPairs1(int[] nums1, int[] nums2, int k) {
        Comparator<List<Integer>> listComparator = Comparator.comparingInt(o -> o.get(0) + o.get(1));
        PriorityQueue<List<Integer>> queue = new PriorityQueue<>(k, listComparator.reversed());
        for (int one : nums1) {
            if (queue.size() == k) {
                List<Integer> peek = queue.peek();
                if ((peek.get(0) + peek.get(1)) <= (one + nums2[0])) {
                    break;
                }
            }
            for (int tow : nums2) {
                if (queue.size() == k) {
                    List<Integer> peek = queue.peek();
                    if ((peek.get(0) + peek.get(1)) <= (one + tow)) {
                        break;
                    }
                    queue.poll();
                }
                ArrayList<Integer> list = new ArrayList<>(2);
                list.add(one);
                list.add(tow);
                queue.add(list);
            }
        }
        return queue.stream().sorted(listComparator).toList();
    }

    // 官方题解
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<>(k, (o1, o2) -> nums1[o1[0]] + nums2[o1[1]] - nums1[o2[0]] - nums2[o2[1]]);
        List<List<Integer>> ans = new ArrayList<>();
        int m = nums1.length;
        int n = nums2.length;
        for (int i = 0; i < Math.min(m, k); i++) {
            pq.offer(new int[]{i, 0});
        }
        while (k-- > 0 && !pq.isEmpty()) {
            int[] idxPair = pq.poll();
            List<Integer> list = new ArrayList<>();
            list.add(nums1[idxPair[0]]);
            list.add(nums2[idxPair[1]]);
            ans.add(list);
            if (idxPair[1] + 1 < n) {
                pq.offer(new int[]{idxPair[0], idxPair[1] + 1});
            }
        }

        return ans;
    }

//    public static void main(String[] args) {
//        Solution KSmallestPairs = new Solution();
//        KSmallestPairs.kSmallestPairs(new int[]{1, 7, 11}, new int[]{2, 4, 6}, 3);
//    }
}

